#include <cassert>

#include "codegen.h"
#include "Type.h"

void CodeGenerater::VisitProgramNode(ProgramAstNode *node) {
    for (auto &func : node->Funcs) {
        func->Accept(this);
    }
}

void CodeGenerater::VisitFunctionNode(FunctionNode *node) {
    CurrentFuncName = node->FuncName;
    printf("\t.text\n");

#ifdef __linux__
    printf("\t.global %s\n", CurrentFuncName.data());
    printf("%s:\n", CurrentFuncName.data());
#else
    // macos
    printf("\t.global _%s\n", CurrentFuncName.data());
    printf("_%s:\n", CurrentFuncName.data());
#endif // __linux__

    int stackSize = 0;
    for (auto &var : node->LocalVars) {
        stackSize += var->Ty->Size;
        stackSize = AlignTo(stackSize, var->Ty->Align);
        var->Offset = -stackSize;
    }

    
    stackSize = AlignTo(stackSize, 16);

    printf("\tpush %%rbp\n");
    printf("\tmov %%rsp, %%rbp\n");
    printf("\tsub $%d, %%rsp\n", stackSize);

    int i = 0;
    for (auto &var : node->Params) {    
        if (var->Ty->Size == 1) {
            printf("\tmov %s, %d(%%rbp)\n", Reg8[i++], var->Offset);
        } else if (var->Ty->Size == 2) {
            printf("\tmov %s, %d(%%rbp)\n", Reg16[i++], var->Offset);
        } else if (var->Ty->Size == 4) {
            printf("\tmov %s, %d(%%rbp)\n", Reg32[i++], var->Offset);
        } else if (var->Ty->Size == 8) {
            printf("\tmov %s, %d(%%rbp)\n", Reg64[i++], var->Offset);
        } else {
          printf("not support type size! got typesize = %d\n", var->Ty->Size);
          assert(0);
        }
        
    }

    for (auto &statm : node->Statms) {
        statm->Accept(this);
        assert(StackTop == 0);
    }
    
    printf(".L.return_%s:\n", CurrentFuncName.data());
    printf("\tmov %%rbp, %%rsp\n");
    printf("\tpop %%rbp\n");
    printf("\tret\n");
}

void CodeGenerater::VisitExprStatmNode(ExprStatmNode *node) {
    if (node->Lhs) {    // 跳过空语句，单分号
        node->Lhs->Accept(this);    
    }
}

void CodeGenerater::VisitFuncCallNode(FuncCallNode *node) {
    for (auto &arg : node->Args) {
        arg->Accept(this);              // 计算函数调用参数值
        Push();                     
    }
    for (int i = node->Args.size() - 1; i >= 0 ; i--) {
        Pop(Reg64[i]);
    }

    std::string funcName(node->FuncName);
#if __linux__
    printf("\tcall %s\n", funcName.data());
#else
    printf("\tcall _%s\n", funcName.data());
#endif
}

void CodeGenerater::VisitDeclarationStatmNode(DeclarationStatmNode *node) {
    for (auto &assignNode : node->AssignNodes) {
        assignNode->Accept(this);
    }
}

void CodeGenerater::VisitIfStatmNode(IfStatmNode *node) {
    int n = Sequece++;
    // 1. 计算条件表达式
    node->Cond->Accept(this);
    // 2. 比较结果
    printf("\tcmp $0, %%rax\n");
    // 3. 存在else语句，生成跳转语句
    if (node->Else) {
        printf("\tje .L.else_%d\n", n);           // 如果rax == 0（cond == false）,跳到else语句   
    } else {
        printf("\tje .L.end_%d\n", n);            // 没有else语句，跳到末尾
    }         
    // 4. 条件真时，执行语句。完毕后，跳过else语句                    
    node->Then->Accept(this);               
    printf("\tjmp .L.end_%d\n", n);   

    // 5. 条件假，执行语句
    if (node->Else) {
        printf(".L.else_%d:\n", n);
        node->Else->Accept(this);
    }
    printf(".L.end_%d:\n", n);
} 

void CodeGenerater::VisitWhileStatmNode(WhileStatmNode *node) {
    int n = Sequece++;
    // 0. while语句开始
    printf(".L.begin_%d:\n", n);
    // 1. 计算条件表达式
    node->Cond->Accept(this);
    // 2. 比较结果
    printf("\tcmp $0, %%rax\n");
    // 3. 条件不满足跳到结尾
    printf("\tje .L.end_%d\n", n);            // 没有else语句，跳到末尾
             
    // 4. 条件真时，执行语句                 
    node->Then->Accept(this);               
    printf("\tjmp .L.begin_%d\n", n);   

    printf(".L.end_%d:\n", n);
} 

void CodeGenerater::VisitForStatmNode(ForStatmNode *node) {
    int n = Sequece++;
    if (node->Init) {
        node->Init->Accept(this);
    }

    printf(".L.begin_%d:\n", n);
    if (node->Cond) {
        node->Cond->Accept(this);           // 1. 计算条件表达式
        printf("\tcmp $0, %%eax\n");        // 2. 比较结果
        printf("\tje .L.end_%d\n", n);      // 3. 条件不满足跳到结尾
    }  // else 条件为空，死循环

    node->Then->Accept(this);           // 4. 条件满足，执行循环体
    if (node->Inc) {
        node->Inc->Accept(this);        // 5. 条件变化
    }
    printf("\tjmp .L.begin_%d\n", n);  
    printf(".L.end_%d:\n", n);
}

void CodeGenerater::VisitReturnStatmNode(ReturnStatmNode *node) {
    node->Lhs->Accept(this);
    printf("\tjmp .L.return_%s\n", CurrentFuncName.data());
}

void CodeGenerater::VisitBlockStatmNode(BlockStatmNode *node) {
    for (auto &statm : node->Statms) {
        statm->Accept(this);
    }
}

void CodeGenerater::VisitAssignExprNode(AssignExprNode *node) {
    GenAddr(node->Lhs.get());                                               // 取局部变量的地址，放入rax
    Push();                                                                 // 保存地址
    node->Rhs->Accept(this);                                                // 获取=号右边值
    SaveRegToMem(node->Ty);                                                     // 值放入地址
}

void CodeGenerater::VisitBinaryNode(BinaryAstNode *node) {
    node->Rhs->Accept(this);
    Push();
    node->Lhs->Accept(this);
    Pop("%rdi");

    switch (node->BinOp) {
    case BinaryOperator::Add: {
        printf("\tadd %%rdi, %%rax\n");
        break;
    }
    case BinaryOperator::Sub: {
        printf("\tsub %%rdi, %%rax\n");
        break;
    }
    case BinaryOperator::Mul: {
        printf("\timul %%rdi, %%rax\n");
        break;
    }
    case BinaryOperator::Div: {
        printf("\tcqo\n");
        printf("\tidiv %%rdi\n");
        break;
    }
    case BinaryOperator::PtrAdd : {
        if (node->Lhs->Ty->IsPointerType()) {
            auto pty = std::dynamic_pointer_cast<PointerType>(node->Lhs->Ty)->Base;
            printf("\timul $%d, %%rdi\n", pty->Size);
        } else {
            auto pty = std::dynamic_pointer_cast<ArrayType>(node->Lhs->Ty)->elementType;
            printf("\timul $%d, %%rdi\n", pty->Size);
        }
        printf("\tadd %%rdi, %%rax\n");
        break;
    }
    case BinaryOperator::PtrSub : {
        if (node->Lhs->Ty->IsPointerType()) {
            auto pty = std::dynamic_pointer_cast<PointerType>(node->Lhs->Ty)->Base;
            printf("\timul $%d, %%rdi\n", pty->Size);
        } else {
            auto pty = std::dynamic_pointer_cast<ArrayType>(node->Lhs->Ty)->elementType;
            printf("\timul $%d, %%rdi\n", pty->Size);
        }
        printf("\tsub %%rdi, %%rax\n");
        break;
    }  
    case BinaryOperator::PtrDiff : {
        printf("\tsub %%rdi, %%rax\n");
        if (node->Lhs->Ty->IsPointerType()) {
            auto pty = std::dynamic_pointer_cast<PointerType>(node->Lhs->Ty)->Base;
            printf("\tmov $%d, %%rdi\n", pty->Size);
        } else {
            auto pty = std::dynamic_pointer_cast<ArrayType>(node->Lhs->Ty)->elementType;
            printf("\tmov $%d, %%rdi\n", pty->Size);
        }
        printf("\tcqo\n");
        printf("\tidiv %%rdi\n");
        break;
    }
    case BinaryOperator::Equal: {
        printf("\tcmp %%rdi, %%rax\n");         // 比较两数是否相等，会设置eflags寄存器
        printf("\tsete %%al\n");                // 使用ZF位设置，相等位1，不等为0
        printf("\tmovzb %%al, %%rax\n");        // 不同大小寄存器传值
        break;
    }
    case BinaryOperator::PipeEqual: {
        printf("\tcmp %%rdi, %%rax\n");         // 比较两数是否相等，会设置eflags寄存器
        printf("\tsetne %%al\n");                // 使用ZF位设置，相等位1，不等为0
        printf("\tmovzb %%al, %%rax\n");        // 不同大小寄存器传值
        break;
    }
    case BinaryOperator::Greater: {
        printf("\tcmp %%rdi, %%rax\n");         // 比较两数是否相等，会设置eflags寄存器
        printf("\tsetg %%al\n");                // 使用ZF位设置，相等位1，不等为0
        printf("\tmovzb %%al, %%rax\n");        // 不同大小寄存器传值
        break;
    }
    case BinaryOperator::GreaterEqual: {
        printf("\tcmp %%rdi, %%rax\n");         // 比较两数是否相等，会设置eflags寄存器
        printf("\tsetge %%al\n");                // 使用ZF位设置，相等位1，不等为0
        printf("\tmovzb %%al, %%rax\n");        // 不同大小寄存器传值
        break;
    }
    case BinaryOperator::Lesser: {
        printf("\tcmp %%rdi, %%rax\n");         // 比较两数是否相等，会设置eflags寄存器
        printf("\tsetl %%al\n");                // 使用ZF位设置，相等位1，不等为0
        printf("\tmovzb %%al, %%rax\n");        // 不同大小寄存器传值
        break;
    }
    case BinaryOperator::LesserEqual: {
        printf("\tcmp %%rdi, %%rax\n");         // 比较两数是否相等，会设置eflags寄存器
        printf("\tsetle %%al\n");                // 使用ZF位设置，相等位1，不等为0
        printf("\tmovzb %%al, %%rax\n");        // 不同大小寄存器传值
        break;
    }
    default:
        assert(0);
        break;
    }
}

void CodeGenerater::VisitUnaryNode(UnaryAstNode *node) {
    switch (node->UnaryOp) {
        case UnaryOperator::Plus: {
            node->Lhs->Accept(this);
            break;
        }
        case UnaryOperator::Minus: {
            node->Lhs->Accept(this);
            printf("\tneg %%rax\n");    // rax结果取反
            break;
        }
        case UnaryOperator::Deref: {
            GenAddr(node);                   // 获取变量地址, 放入rax
            LoadMemToReg(node->Ty);                         // 从地址取值，放入rax
            break;
        }
        case UnaryOperator::Amp: {
            GenAddr(node->Lhs.get());
            break;
        }
    }
}

void CodeGenerater::VisitConstantNode(ConstantAstNode *node) {
    printf("\tmov $%d, %%rax\n", node->Value);
}

void CodeGenerater::VisitSizeofNode(SizeofNode *node) {
    printf("\tmov $%d, %%rax\n", node->Lhs->Ty->Size);
}

void CodeGenerater::VisitVarNode(VarAstNode *node) {
    GenAddr(node);                                       // 获取变量地址, 放入rax
    LoadMemToReg(node->Ty);                              // 从地址取值 放入rax,如果是数组就不用取值
}

void CodeGenerater::Push() {
    printf("\tpush %%rax\n");
    StackTop++;
}

void CodeGenerater::Pop(const char* reg) {
    printf("\tpop %s\n", reg);
    StackTop--;
}

void CodeGenerater::LoadMemToReg(std::shared_ptr<Type> ty) {
    if (ty->IsArrayType()) {
        return ;
    }
    if (ty->Size == 1) {
        printf("\tmovsb (%%rax), %%rax\n");
    } else if (ty->Size == 2) {
        printf("\tmovsw (%%rax), %%rax\n");
    } else if (ty->Size == 4) {
        printf("\tmovsl (%%rax), %%rax\n");
    } else if (ty->Size == 8) {
        printf("\tmov (%%rax), %%rax\n");
    } else {
        printf("not support type size! got typesize = %d\n", ty->Size);
        assert(0);
    }
}

void CodeGenerater::SaveRegToMem(std::shared_ptr<Type> ty) {
    Pop("%rdi");
    if (ty->Size == 1) {
        printf("\tmov %%al, (%%rdi)\n");
    } else if (ty->Size == 2) {
        printf("\tmov %%ax, (%%rdi)\n");
    } else if (ty->Size == 4) {
        printf("\tmov %%eax, (%%rdi)\n");
    } else if (ty->Size == 8) {
        printf("\tmov %%rax, (%%rdi)\n");
    } else {
        printf("not support type size! got typesize = %d\n", ty->Size);
        assert(0);
    }
}


void CodeGenerater::GenAddr(AstNode *node) {
    if (auto var = dynamic_cast<VarAstNode *>(node)) {
        printf("\tlea %d(%%rbp), %%rax\n", var->VarObj->Offset);        // 取地址符后跟变量节点
    } else if (auto unaryNode = dynamic_cast<UnaryAstNode *>(node)) {
        if (unaryNode->UnaryOp == UnaryOperator::Deref) {               // 解引用后跟变量
            unaryNode->Lhs->Accept(this);
        } else {
            printf("this unarynode must be defer!\n");
            assert(0);
        }
    } else {
        printf("not a lvalue!\n");
        assert(0);
    }
}

int CodeGenerater::AlignTo(int size, int align) {
    return (size + align - 1) / align * align;
}
